#pragma once

#include "ast.h"
#include "location.h"
#include "pass.h"

#define REGISTER_INTERFACE(className) \
    virtual AST::ASTNode* Visit##className(AST::className* node) = 0;
#define REGISTER_VISIT(className) \
    if (auto* nxt = node->As<AST::className>()) ret = Visit##className(nxt);

class ASTVisitor : public Pass {
    Location curLoc;

  protected:
    REGISTER_INTERFACE(AssignExpr)
    REGISTER_INTERFACE(CallExpr)
    REGISTER_INTERFACE(ArrayExpr)
    REGISTER_INTERFACE(OperatorExpr)
    REGISTER_INTERFACE(TernaryExpr)
    REGISTER_INTERFACE(LiteralExpr)
    REGISTER_INTERFACE(VoidExpr)
    REGISTER_INTERFACE(IdentifierExpr)
    REGISTER_INTERFACE(InitializerExpr)
    REGISTER_INTERFACE(CastExpr)
    REGISTER_INTERFACE(AddrExpr)
    REGISTER_INTERFACE(ArgList)
    REGISTER_INTERFACE(Block)
    REGISTER_INTERFACE(VarDecl)
    REGISTER_INTERFACE(FuncDecl)
    REGISTER_INTERFACE(IfStmt)
    REGISTER_INTERFACE(WhileStmt)
    REGISTER_INTERFACE(ContinueStmt)
    REGISTER_INTERFACE(BreakStmt)
    REGISTER_INTERFACE(ReturnStmt)

  public:
    virtual void Visit(AST::ASTNode*& node) {
        AST::ASTNode* ret = nullptr;
        curLoc = node->loc;
        REGISTER_VISIT(AssignExpr)
        REGISTER_VISIT(CallExpr)
        REGISTER_VISIT(ArrayExpr)
        REGISTER_VISIT(OperatorExpr)
        REGISTER_VISIT(TernaryExpr)
        REGISTER_VISIT(LiteralExpr)
        REGISTER_VISIT(VoidExpr)
        REGISTER_VISIT(IdentifierExpr)
        REGISTER_VISIT(InitializerExpr)
        REGISTER_VISIT(CastExpr)
        REGISTER_VISIT(AddrExpr)
        REGISTER_VISIT(ArgList)
        REGISTER_VISIT(Block)
        REGISTER_VISIT(VarDecl)
        REGISTER_VISIT(FuncDecl)
        REGISTER_VISIT(IfStmt)
        REGISTER_VISIT(WhileStmt)
        REGISTER_VISIT(ContinueStmt)
        REGISTER_VISIT(BreakStmt)
        REGISTER_VISIT(ReturnStmt)
        node = ret;
    }

    virtual void* Run(void* input) override {
        auto node = (AST::ASTNode*)input;
        Visit(node);
        return node;
    }
};

#undef REGISTER_INTERFACE
#undef REGISTER_VISIT
